#include "clang/Frontend/FrontendPluginRegistry.h"
#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/AST/ParentMap.h"
#include "clang/AST/ParentMapContext.h"
#include "clang/Sema/Sema.h"
#include "llvm/Support/raw_ostream.h"

using namespace clang;

namespace {

static std::unordered_map<std::string, std::string> argsMap;

class ClangOjbCCodeSpecCheckerVisitor: public RecursiveASTVisitor<ClangOjbCCodeSpecCheckerVisitor> {
private:
    inline int argsCount() {
        return argsMap.count("ArgsCount") ? std::stoi(argsMap["ArgsCount"]) : 5;
    }
    
    inline int maxLineCount() {
        return argsMap.count("MaxLineCount") ? std::stoi(argsMap["MaxLineCount"]) : 100;
    }
    
    bool ObjCPropIvarNameCamelCased(clang::StringRef& name, bool ivar = false) {
        if (ivar) {
            return name.starts_with("_") && ::islower(*(name.data() + 1));
        }
        
        return ::islower(*name.data());
    }
    
    bool OBJCNamePascalCased(clang::StringRef& name) {
        return ::isupper(*name.data());
    }
    
    bool ObjCClassEnumNamePascalCased(clang::StringRef& name) {
        if (!argsMap.count("ClsPrefix")) {
            return true;
        }
        
        return name.starts_with(argsMap["ClsPrefix"]) && ::isupper(*(name.data() + 2));
    }
    
    bool isUserSourceFile(const clang::SourceLocation& loc) {
        clang::StringRef fileName = this->Context->getSourceManager().getFilename(loc);
        
        if (SM.isInSystemHeader(loc)) {
            return false;
        }
        
        return (int)fileName.find("/Applications/Xcode.app/") == -1;
    }
    
    bool isTypeNullability(const clang::QualType& type) {
        clang::QualType desugared = type.getNonReferenceType().getUnqualifiedType();
        
        if (const ObjCObjectPointerType *objcPtrType = desugared->getAs<clang::ObjCObjectPointerType>()) {
            std::optional<NullabilityKind> nullable = objcPtrType->getNullability();
            if (nullable.has_value() && nullable.value() == NullabilityKind::Nullable) {
                return true;
            }
        }
        
        if (const clang::PointerType *ptrType = desugared->getAs<clang::PointerType>()) {
            std::optional<NullabilityKind> nullable = ptrType->getNullability();
            if (nullable.has_value() && nullable.value() == NullabilityKind::Nullable) {
                return true;
            }
        }
        
        return false;
    }
    
    bool findOnlyNewLine(const char *txt, bool forward = true) {
        if (!txt) {
            return false;
        }
        
        int idx = 0;
        while (*(txt + idx) != '\n') {
            if (*(txt + idx) == '\t' || *(txt + idx) == ' ') {
                idx += forward ? 1 : -1;
                continue;
            }
            
            return false;
        }
        
        return true;
    }
    
    void ObjCCompoundValid(const SourceLocation& beg, bool *vBeg, const SourceLocation& end, bool *vEnd) {
        const char *st = SM.getCharacterData(beg);
        const char *ed = SM.getCharacterData(end);
        
        
        if (st && (*(st + 1) != ' ' || *(st + 2) != '{')) {
            *vBeg = false;
        }
        
        if (!ed || *ed != '}') {
            return ;
        }
        
        *vEnd = this->findOnlyNewLine(ed - 1, false);
    }
    
    bool ObjCPropertyAttributesOrder(llvm::SmallVector<clang::StringRef>& attrs) {
        int attrIndex = -1;
        
        auto cmp = [&attrIndex](int index) -> bool {
            if (index < attrIndex) {
                return false;
            }
            
            attrIndex = index;
            return true;
        };
        
        for (clang::StringRef &s : attrs) {
            int curIndex = -1;
            clang::StringRef str = s.trim();
            
            if (!str.compare("atomic") || !str.compare("nonatomic")) {
                curIndex = 0;
            } else if (!str.compare("assign") || !str.compare("strong") || !str.compare("weak") || !str.compare("copy") || !str.compare("retain") || !str.compare("unsafe_unretained")) {
                curIndex = 1;
            } else if (!str.compare("readonly") || !str.compare("readwrite")) {
                curIndex = 2;
            } else if (!str.compare("nullable") || !str.compare("nonnull")) {
                curIndex = 3;
            } else if (!str.compare("setter") || !str.compare("getter")) {
                curIndex = 4;
            } else if (!str.compare("class")) {
                curIndex = 5;
            }
            
            if (!cmp(curIndex)) {
                return false;
            }
        }
        
        return true;
    }
    
    // 辅助函数：判断子语句是否在父语句的子树中
    bool isStmtInSubtree(const Stmt *ParentStmt, const Stmt *ChildStmt) {
        if (!ParentStmt || !ChildStmt)
            return false;
        
        // 如果父语句是复合语句，检查子语句是否在其中
        if (const CompoundStmt *CS = dyn_cast<CompoundStmt>(ParentStmt)) {
            for (Stmt *S : CS->body()) {
                if (S == ChildStmt || isStmtInSubtree(S, ChildStmt)) {
                    return true;
                }
            }
        } else {
            // 非复合语句（如单个表达式），直接比较
            return ParentStmt == ChildStmt;
        }
        
        return false;
    }
    
    const Expr *getBlockExprFromCall(const CallExpr *CE) {
        const clang::Expr *Callee = CE->getCallee()->IgnoreParenCasts();
        
        if (const PseudoObjectExpr *pExpr = dyn_cast<PseudoObjectExpr>(Callee)) {
            if (pExpr->getType()->isBlockPointerType()) {
                return pExpr;
            }
        }
        
        if (const DeclRefExpr *rExpr = dyn_cast<clang::DeclRefExpr>(Callee)) {
            if (rExpr->getType()->isBlockPointerType()) {
                return rExpr;
            }
        }
        
        if (const clang::ObjCIvarRefExpr *rExpr = dyn_cast<clang::ObjCIvarRefExpr>(Callee)) {
            if (rExpr->getType()->isBlockPointerType()) {
                return rExpr;
            }
        }
        
        return nullptr;
    }
    
    bool isNullExpr(const clang::Expr *E) {
        E = E->IgnoreImpCasts()->IgnoreParenCasts();
        if (const clang::IntegerLiteral *pExpr = dyn_cast<clang::IntegerLiteral>(E)) {
            if (pExpr->getValue().getSExtValue() == 0) {
                return true;
            }
        }
        
        return false;
    }
    
    std::string getBlockVarDeclName(const clang::Expr *E) {
        if (const PseudoObjectExpr *tExpr = dyn_cast<PseudoObjectExpr>(E)) {
            for (auto it = tExpr->child_begin(); it != tExpr->child_end(); ++it) {
                if (const clang::ObjCPropertyRefExpr *ex = dyn_cast<clang::ObjCPropertyRefExpr>(*it)) {
                    return ex->getExplicitProperty()->getPropertyIvarDecl()->getNameAsString();
                }
            }
        } else if (const clang::ObjCIvarRefExpr *rExpr = dyn_cast<clang::ObjCIvarRefExpr>(E)) {
            return rExpr->getDecl()->getNameAsString();
        } else if (const clang::DeclRefExpr *rExpr = dyn_cast<clang::DeclRefExpr>(E)) {
            return rExpr->getDecl()->getNameAsString();
        }
        
        return "";
    }

    bool isBlockNullCheck(const clang::Expr *Cond, const clang::Expr *CallExpr) {
        if (!Cond || !CallExpr) {
            return false;
        }
        
        Cond = Cond->IgnoreParenCasts();
        CallExpr = CallExpr->IgnoreParenCasts();
        
        if (const clang::BinaryOperator *bOpExpr = dyn_cast<clang::BinaryOperator>(Cond)) {
            /// if (self.block != nil)
            /// if (_block != nil)
            /// if (block != nil)
            
            const clang::Expr *lhs = bOpExpr->getLHS()->IgnoreParenImpCasts();
            const clang::Expr *rhs = bOpExpr->getRHS()->IgnoreParenImpCasts();
            
            if (!lhs || !rhs) {
                return false;
            }
            
            if (this->isNullExpr(lhs)) {
                std::swap(lhs, rhs);
            }
            
            if (!this->isNullExpr(rhs)) {
                return false;
            }
            
            std::string CallExprName = this->getBlockVarDeclName(CallExpr);
            std::string BlockName = this->getBlockVarDeclName(lhs);
            
            return !CallExprName.compare(BlockName);
        } else {
            /// if (self.block)
            /// if (_block)
            /// if (block)
            return !this->getBlockVarDeclName(Cond).compare(this->getBlockVarDeclName(CallExpr));
        }
        
        return false;
    }
    
    clang::DiagnosticsEngine::Level diagnosticsLevel(void) {
        if (!argsMap.count("DiagLevel")) {
            return clang::DiagnosticsEngine::Warning;
        }

        std::string level = argsMap["DiagLevel"];
        if (!level.compare("Ignored")) {
            return clang::DiagnosticsEngine::Ignored;
        } else if (!level.compare("Note")) {
            return clang::DiagnosticsEngine::Note;
        } else if (!level.compare("Remark")) {
            return clang::DiagnosticsEngine::Remark;
        } else if (!level.compare("Error")) {
            return clang::DiagnosticsEngine::Error;
        } else if (!level.compare("Fatal")) {
            return clang::DiagnosticsEngine::Fatal;
        }
        
        return clang::DiagnosticsEngine::Warning;
    }
    
private:
    void handleReturnStmt(const clang::ReturnStmt *S) {
        if (!this->curMethod) {
            return ;
        }
        
        const clang::Expr *expr = S->getRetValue()->IgnoreParenImpCasts();
        if (!expr) {
            return ;
        }
        
        clang::StringRef retType = clang::StringRef(this->curMethod->getReturnType().getAsString());
        
        std::string retValTypeStr = std::string();
        bool isNilRetVal = false;
        
        if (const clang::DeclRefExpr *refExpr = dyn_cast<clang::DeclRefExpr>(expr)) {
            if (const clang::VarDecl *varDecl = dyn_cast<clang::VarDecl>(refExpr->getDecl())) {
                retValTypeStr = varDecl->getType().getAsString();
            }
        } else if (const clang::ObjCMessageExpr *msgExpr = dyn_cast<clang::ObjCMessageExpr>(expr)) {
            retValTypeStr = msgExpr->getType().getAsString();
        } else if (this->isNullExpr(expr)) {
            isNilRetVal = true;
        }
        
        /// 返回nil指针，但返回类型未声明nullable
        {
            if (isNilRetVal) {
                llvm::outs() << this->curMethod->getSelector().getAsString();
                
                bool retTypeNullable = this->isTypeNullability(this->curMethod->getReturnType());
                
                if (!retTypeNullable) {
                    DE.Report(expr->getExprLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】函数返回值类型:「%0」, 实际返回为「nil」, 请指定「Nullable」")) << retType;
                }
            }
        }
        
        if (retValTypeStr.empty()) {
            return ;
        }
        
        /// 不可变容器返回可变容器类型
        {
            clang::StringRef retValType = clang::StringRef(retValTypeStr);
            bool retMulToIMul = (retType.starts_with("NSArray") && retValType.starts_with("NSMutableArray")) || (retType.starts_with("NSDictionary") && retValType.starts_with("NSMutableDictionary")) || (retType.starts_with("NSString") && retValType.starts_with("NSMutableString")) || (retType.starts_with("NSAttributedString") && retValType.starts_with("NSMutableAttributedString")) || (retType.starts_with("NSData") && retValType.starts_with("NSMutableData")) || (retType.starts_with("NSSet") && retValType.starts_with("NSMutableSet"));
            
            if (retMulToIMul) {
                DE.Report(expr->getExprLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】函数返回值类型:「%0」, 实际返回类型「%1」, 请使用「copy」")) << retType << retValType;
            }
        }
    }
    
    void handleIfStmt(const clang::IfStmt* S) {
        if (!this->curMethod) {
            return ;
        }
        
        clang::SourceLocation rParLoc = S->getRParenLoc();
        
        const Stmt *then = S->getThen();
        if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(then)) {
            SourceLocation loc = cStmt->getEndLoc();
            
            bool vBeg = true, vEnd = true;
            this->ObjCCompoundValid(rParLoc, &vBeg, loc, &vEnd);
            
            if (!vBeg) {
                DE.Report(rParLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应该与「if语句」在「同一行」，并且中间应有「一个空格」"));
            }
            
            if (!vEnd) {
                DE.Report(loc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「if语句」的「}」 应该「单独占一行」"));
            }
            
        } else {
            DE.Report(rParLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「if」语句应该添加「{}」, 并且「{」应与「if语句」在「同一行」，中间应有「一个空格」"));
        }
        
        const Stmt *els = S->getElse();
        if (els) {
            if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(els)) {
                SourceLocation begLoc = cStmt->getBeginLoc();
                SourceLocation endLoc = cStmt->getEndLoc();
                
                const char *begTxt = SM.getCharacterData(begLoc);
                const char *endTxt = SM.getCharacterData(endLoc);
                
                bool vBeg = this->findOnlyNewLine(begTxt, false), vEnd = this->findOnlyNewLine(endTxt, false);
                
                if (vBeg) {
                    DE.Report(begLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「else」语句的「{」应该与「else」语句在「同一行」，并且中间应有「一个空格」"));
                }
                
                if (!vEnd) {
                    DE.Report(endLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「else」语句的「}」 应该「单独占一行」"));
                }
            } else {
                DE.Report(els->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「else」语句应该添加「{}」, 并且「{」应与「else」语句在「同一行」，中间应有「一个空格」"));
            }
        }
    }
    
    void handleCompoundStmt(const clang::CompoundStmt *S) {
        if (!this->curMethod) {
            return ;
        }
        
        /// 方法行数检查
        {
            clang::SourceLocation begLoc = S->getBeginLoc();
            clang::SourceLocation endLoc = S->getEndLoc();
            
            unsigned int begLineNumber = SM.getLineNumber(SM.getFileID(begLoc), SM.getFileOffset(begLoc));
            unsigned int endLineNumber = SM.getLineNumber(SM.getFileID(endLoc), SM.getFileOffset(endLoc));
            
            if (endLineNumber - begLineNumber + 1 > this->maxLineCount()) {
                DE.Report(begLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「方法行数」最多不能超过「%0行」")) << this->maxLineCount();
            }
        }
    }
    
    void handleForStmt(const clang::ForStmt *S) {
        const Stmt *st = S->getBody();
        if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(st)) {
            
            SourceLocation bLoc = cStmt->getBeginLoc();
            SourceLocation eLoc = cStmt->getEndLoc();
            
            {
                bool vBeg = true, vEnd = true;
                this->ObjCCompoundValid(bLoc, &vBeg, eLoc, &vEnd);
                
                if (!vBeg) {
                    DE.Report(bLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应该与「for语句」在「同一行」，并且中间应有「一个空格」"));
                }
                
                if (!vEnd) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「for语句」的「}」 应该「单独占一行」"));
                }
            }
            
            {
                const char *btxt = SM.getCharacterData(bLoc);
                const char *etxt = SM.getCharacterData(eLoc);
                
                bool newLine = this->findOnlyNewLine(btxt - 1, false);
                if (newLine || *(btxt - 1) != ' ') {
                    DE.Report(cStmt->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与「for」在「同一行」，且中间应有「一个空格」"));
                }
                
                bool next = this->findOnlyNewLine(etxt - 1, false);
                if (!next) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「}」应「另起一行」"));
                }
            }
        } else {
            DE.Report(st->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「for」语句应该使用「{}」"));
        }
    }
    
    void handleWhileStmt(const clang::WhileStmt *S) {
        const Stmt *st = S->getBody();
        if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(st)) {
            
            SourceLocation bLoc = cStmt->getBeginLoc();
            SourceLocation eLoc = cStmt->getEndLoc();
            
            {
                bool vBeg = true, vEnd = true;
                
                this->ObjCCompoundValid(bLoc, &vBeg, eLoc, &vEnd);
                if (!vBeg) {
                    DE.Report(bLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应该与「while语句」在「同一行」，并且中间应有「一个空格」"));
                }
                
                if (!vEnd) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「while语句」的「}」 应该「单独占一行」"));
                }
            }
            
            {
                const char *btxt = SM.getCharacterData(bLoc);
                const char *etxt = SM.getCharacterData(eLoc);
                
                bool newLine = this->findOnlyNewLine(btxt - 1, false);
                if (newLine || *(btxt - 1) != ' ') {
                    DE.Report(bLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与「while」在同一行，且中间应有「一个空格」"));
                }
                
                bool next = this->findOnlyNewLine(etxt - 1, false);
                if (!next) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「}」应「另起一行」"));
                }
            }
        }
    }
    
    void handleDoStmt(const clang::DoStmt *S) {
        const Stmt *st = S->getBody();
        if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(st)) {
            
            SourceLocation bLoc = cStmt->getBeginLoc();
            SourceLocation eLoc = cStmt->getEndLoc();
            
            {
                bool vBeg = true, vEnd = true;
                this->ObjCCompoundValid(bLoc, &vBeg, eLoc, &vEnd);
                
                if (!vBeg) {
                    DE.Report(bLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应该与「do语句」在「同一行」，并且中间应有「一个空格」"));
                }
                
                if (!vEnd) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「do语句」的「}」 应该「单独占一行」"));
                }
            }
            
            {
                const char *btxt = SM.getCharacterData(bLoc);
                const char *etxt = SM.getCharacterData(eLoc);
                
                bool newLine = this->findOnlyNewLine(btxt - 1, false);
                if (newLine || *(btxt - 1) != ' ') {
                    DE.Report(S->getDoLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与「do」在「同一行」，且中间应有「一个空格」"));
                }
                
                bool next = this->findOnlyNewLine(etxt + 1);
                if (next) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「while」语句应与「}」在「同一行」"));
                }
            }
        }
    }
    
    void handleSwitchStmt(const clang::SwitchStmt *S) {
        const Stmt *stmt = S->getBody();
        if (const CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(stmt)) {
            SourceLocation bLoc = cStmt->getBeginLoc();
            SourceLocation eLoc = cStmt->getEndLoc();
            
            {
                const char *btxt = SM.getCharacterData(bLoc);
                const char *etxt = SM.getCharacterData(eLoc);
                
                bool newLine = this->findOnlyNewLine(btxt - 1, false);
                if (newLine || *(btxt - 1) != ' ') {
                    DE.Report(bLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与「switch」在「同一行」，且中间应有「一个空格」"));
                }
                
                bool next = this->findOnlyNewLine(etxt + 1);
                if (!next) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「switch」语句的「}」应「单独占一行」"));
                }
            }
        }
    }
    
    void handleCaseStmt(const clang::CaseStmt *S) {
        const Stmt *stmt = S->getSubStmt();
        if (const CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(stmt)) {
            SourceLocation bLoc = cStmt->getBeginLoc();
            SourceLocation eLoc = cStmt->getEndLoc();
            
            {
                const char *btxt = SM.getCharacterData(bLoc);
                const char *etxt = SM.getCharacterData(eLoc);
                
                bool newLine = this->findOnlyNewLine(btxt - 1, false);
                if (newLine || *(btxt - 1) != ' ') {
                    DE.Report(bLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与「case」在「同一行」，且中间应有「一个空格」"));
                }
                
                bool next = this->findOnlyNewLine(etxt + 1);
                if (!next) {
                    DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「case」语句的「}」应「单独占一行」"));
                }
            }
        } else {
            DE.Report(S->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「case」语句应该使用「{}」"));
        }
    }
    
    void handleBreakStmt(const BreakStmt *S) {
        SourceLocation bLoc = S->getBeginLoc();
        SourceLocation eLoc = S->getEndLoc();
        
        {
            const char *btxt = SM.getCharacterData(bLoc);
            const char *etxt = SM.getCharacterData(eLoc);
            
            bool newLine = this->findOnlyNewLine(btxt - 1, false);
            bool next = this->findOnlyNewLine(etxt + 1);
            
            if (!newLine || !next) {
                DE.Report(eLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「break」语句应「单独占一行」"));
            }
        }
    }
    
public:
    explicit ClangOjbCCodeSpecCheckerVisitor(clang::CompilerInstance& CI)
    : CI(CI), Context(&CI.getASTContext()), SM(CI.getSourceManager()), DE( CI.getDiagnostics()), PM(nullptr), curMethod(nullptr) {}
    
    bool TraverseDecl(Decl *D) {
        if (D->getBody()) {
            this->PM = std::make_unique<clang::ParentMap>(D->getBody());
        }
        return this->RecursiveASTVisitor<ClangOjbCCodeSpecCheckerVisitor>::TraverseDecl(D);
    }
    
    bool VisitObjCInterfaceDecl(const ObjCInterfaceDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        
        /// 类名
        {
            clang::StringRef clsName = D->getName();
            if (!this->ObjCClassEnumNamePascalCased(clsName)) {
                DE.Report(D->getLocation(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「%0」应该以「%1」开头，并遵循「大驼峰命名法」")) << clsName << argsMap["ClsPrefix"];
            }
        }
        
        return true;
    }
    
    bool VisitBinaryOperator(const clang::BinaryOperator *E) {
        if (!this->isUserSourceFile(E->getBeginLoc())) {
            return true;
        }
        
        
        clang::SourceLocation opCodeLoc = E->getOperatorLoc();
        const char *opTxt = SM.getCharacterData(opCodeLoc);
        
        if (*(opTxt - 1) != ' ') {
            DE.Report(opCodeLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「左表达式」与「操作符」之间应有「一个空格」"));
        }
        
        clang::StringRef opCodeStr = E->getOpcodeStr();
        if (*(opTxt + opCodeStr.size()) != ' ') {
            DE.Report(opCodeLoc.getLocWithOffset(opCodeStr.size()), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「右表达式」与「操作符」之间应有「一个空格」"));
        }
        
        return true;
    }
    
    bool VisitUnaryOperator(const clang::UnaryOperator *E) {
        if (!this->isUserSourceFile(E->getBeginLoc())) {
            return true;
        }
        
        clang::SourceLocation opCodeLoc = E->getOperatorLoc();
        const char *opTxt = SM.getCharacterData(opCodeLoc);
        
        clang::StringRef opCodeStr = E->getOpcodeStr(E->getOpcode());
        
        if (E->getOpcode() == clang::UnaryOperatorKind::UO_PostInc || E->getOpcode() == clang::UnaryOperatorKind::UO_PostDec) {
            if (*(opTxt + opCodeStr.size()) != ' ') {
                DE.Report(opCodeLoc.getLocWithOffset(-1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「左表达式」与「操作符」之间不应有「空格」"));
            }
        } else {
            if (*(opTxt - 1) != ' ') {
                DE.Report(opCodeLoc.getLocWithOffset(opCodeStr.size()), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「右表达式」与「操作符」之间不应有「个空格」"));
            }
        }
        
        return true;
    }
    
    bool VisitCallExpr(const clang::CallExpr *E) {
        if (!this->isUserSourceFile(E->getBeginLoc()) || !this->curMethod) {
            return true;
        }
        
        /// block判空检查
        {
            const clang::Expr *cExpr = this->getBlockExprFromCall(E);
            if (cExpr) {
                if (const CompoundStmt *cStmt = dyn_cast<CompoundStmt>(this->curMethod->getBody())) {
                    bool hasValidCheck = false, hasMultiComp = false;
                    for (auto it = cStmt->child_begin(); it != cStmt->child_end(); ++it) {
                        if (const clang::IfStmt *ifStmt = dyn_cast<clang::IfStmt>(*it)) {
                            const Expr *cond = ifStmt->getCond();
                            const Stmt *then = ifStmt->getThen();
                            
                            if (this->isStmtInSubtree(then, E)) {
                                if (const clang::BinaryOperator *bOpExpr =
                                    dyn_cast<clang::BinaryOperator>(cond)) {
                                    if (bOpExpr->getOpcode() != clang::BO_NE) {
                                        hasMultiComp = true;
                                        break;
                                    }
                                }
                                
                                if (this->isBlockNullCheck(cond, cExpr)) {
                                    hasValidCheck = true;
                                    break;
                                }
                            }
                        }
                    }
                    
                    if (!hasValidCheck) {
                        if (hasMultiComp) {
                            DE.Report(E->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「if」有多个检查条件，为提高代码可读性，请确保「Block」只有「一个」检查条件"));
                        } else {
                            DE.Report(E->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】当前「Block」的调用未检测到「if判断」，可能引起Crash问题，请检查"));
                        }
                    }
                }
            }
        }
        
        /// 参数格式检查
        {
            bool hasError = false;
            
            unsigned int argCnt = E->getNumArgs();
            if (argCnt) {
                for (unsigned int i = 0; i < argCnt - 1; ++i) {
                    const Expr *aExpr = E->getArg(i);
                    
                    SourceLocation begLoc = aExpr->getBeginLoc();
                    SourceLocation endLoc = aExpr->getEndLoc();
                    
                    const char *begTxt = SM.getCharacterData(begLoc);
                    const char *endTxt = SM.getCharacterData(endLoc);
                    
                    if (i == 0 && *(begTxt - 1) != '(') {
                        hasError = true;
                        break;
                    }
                    
                    if (*(endTxt + 1) != ',') {
                        hasError = true;
                        break;
                    } else if (*(endTxt + 1) == ',' && *(endTxt + 2) != ' ') {
                        hasError = true;
                        break;
                    }
                }
                
                if (hasError) {
                    DE.Report(E->getArg(0)->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】参数列表「第一个参数」应该跟随在「(」后面，且「随后参数」的「,」应「跟随前一个参数」，并与后一个参数中间保持「一个空格」"));
                }
            }
        }
        
        return true;
    }
    
    bool VisitObjCMethodDecl(const ObjCMethodDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        
        this->curMethod = D;
        
        clang::StringRef mName = clang::StringRef(D->getSelector().getAsString());
        clang::StringRef retType =clang::StringRef(D->getReturnType().getAsString());
        
        {
            size_t argSize = D->param_size();
            /// 方法参数个数检查
            {
                if (argSize > this->argsCount()) {
                    DE.Report(D->getLocation(), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「%0」的参数个数不能超过「%1」个")) << mName << this->argsCount();
                }
            }
            
            /// 方法参数检查
            {
                for (unsigned int i = 0; i != argSize; ++i) {
                    const clang::ParmVarDecl *param = D->getParamDecl(i);
                    if (!param) {
                        continue;
                    }
                    
                    clang::QualType type = param->getType();
                    if (type->isObjCObjectPointerType()) {
                        clang::SourceLocation endLoc = param->getTypeSpecEndLoc();
                        
                        const char *endTxt = SM.getCharacterData(endLoc);
                        if (*(endTxt - 1) != ' ') {
                            DE.Report(endLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与类型之间应有「一个空格」"));
                        }
                    }
                }
            }
        }
        
        /// 初始化器必须返回instancetype
        {
            if (mName.starts_with("init") && retType.compare("instancetype")) {
                DE.Report(D->getSelectorLoc(0), DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「%0」的返回类型必须为「instancetype」")) << mName;
            }
        }
        
        /// 方法声明 -/+ 与方法名中间添加空格
        {
            SourceLocation loc = D->getBeginLoc();
            const char *txt = SM.getCharacterData(loc);

            if (*(txt + 1) != ' ') {
                DE.Report(loc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「%0」与方法声明中间应有「一个空格」")) << (D->isClassMethod() ? "+" : "-");
            }
        }
        
        /// 左大括号与方法名同行，右大括号换行
        {
            const Stmt *body = D->getBody();
            if (!body) {
                return true;
            }
            
            if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(body)) {
                SourceLocation begLoc = cStmt->getBeginLoc();
                SourceLocation endLoc = cStmt->getEndLoc();
                
                const char *begTxt = SM.getCharacterData(begLoc);
                const char *endTxt = SM.getCharacterData(endLoc);
                
                bool vBeg = this->findOnlyNewLine(begTxt - 1, false), vEnd = this->findOnlyNewLine(endTxt - 1, false);
                
                if (vBeg) {
                    DE.Report(begLoc.getLocWithOffset(1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「{」应与方法声明「%0」在同一行，并且中间应有「一个空格」")) << mName;
                }
                
                if (!vEnd) {
                    DE.Report(endLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「}」 应单独占一行"));
                }
            }
        }
        
        return true;
    }
    
    bool VisitStmt(const Stmt* S) {
        SourceLocation loc = S->getBeginLoc();
        if (!this->isUserSourceFile(loc)) {
            return true;
        }
        
        if (const clang::ReturnStmt *retStmt = dyn_cast<clang::ReturnStmt>(S)) {
            handleReturnStmt(retStmt);
        } else if (const clang::IfStmt *ifStmt = dyn_cast<clang::IfStmt>(S)) {
            handleIfStmt(ifStmt);
        } else if (const clang::ForStmt *forStmt = dyn_cast<clang::ForStmt>(S)) {
            handleForStmt(forStmt);
        } else if (const clang::WhileStmt *whileStmt = dyn_cast<clang::WhileStmt>(S)) {
            handleWhileStmt(whileStmt);
        } else if (const clang::BreakStmt *breakStmt = dyn_cast<clang::BreakStmt>(S)) {
            handleBreakStmt(breakStmt);
        } else if (const clang::DoStmt *doStmt = dyn_cast<clang::DoStmt>(S)) {
            handleDoStmt(doStmt);
        } else if (const clang::SwitchStmt *sStmt = dyn_cast<clang::SwitchStmt>(S)) {
            handleSwitchStmt(sStmt);
        } else if (const clang::CaseStmt *cStmt = dyn_cast<clang::CaseStmt>(S)) {
            handleCaseStmt(cStmt);
        } else if (const clang::CompoundStmt *cStmt = dyn_cast<clang::CompoundStmt>(S)) {
            handleCompoundStmt(cStmt);
        }
        
        return true;
    }
    
    bool VisitEnumDecl(const EnumDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        
        SourceLocation loc = D->getLocation();
        if (SM.isInSystemHeader(loc) || !D->isComplete() || D->enumerators().empty()) {
            return true;
        }
        
        QualType underlyingType = D->getIntegerType();
        
        if (underlyingType.getAsString().compare("NSInteger") && underlyingType.getAsString().compare("NSUInteger")) {
            DE.Report(loc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「枚举类型」请使用「NSInteger/NSUInteger」"));
        }
        
        clang::StringRef enumName = D->getName();
        if (!this->ObjCClassEnumNamePascalCased(enumName)) {
            DE.Report(loc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「枚举名」请以「%0」开头，并遵循「大驼峰命名法」")) << argsMap["ClsPrefix"];
        }
        
        for (auto it = D->enumerator_begin(); it != D->enumerator_end(); ++it) {
            if (!it->getName().starts_with(enumName)) {
                DE.Report(it->getLocation(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「枚举项」请以「%0」开头，并遵循「大驼峰命名法」")) << enumName;
            }
        }
        
        return true;
    }
    
    bool VisitVarDecl(const VarDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        
        clang::QualType type = D->getType();
        
        {
            if (type->isObjCObjectPointerType() || type->isPointerType()) {
                clang::SourceLocation endLoc = D->getEndLoc();
                const char *txt = SM.getCharacterData(endLoc);
                
                if (*(txt - 1) != ' ') {
                    std::string typeName = type.getAsString();
                    DE.Report(endLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间至少要有「一个空格」")) << clang::StringRef(typeName.substr(0, typeName.size() - 1)).trim();
                }
                
                clang::StringRef varName = D->getName();
                if (varName.size() > 0 && *(txt - 1) == ' ') {
                    DE.Report(endLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间不应该有「任何空格」")) << varName;
                }
            }
        }
        
        return true;
    }
    
    bool VisitObjCIvarDecl(const ObjCIvarDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        clang::LangOptions LangOpts = this->Context->getLangOpts();
        
        clang::StringRef ivarType = clang::StringRef(D->getType().getAsString());
        clang::StringRef ivarName = clang::StringRef(D->getNameAsString());
        
        /// 1. 检查\*与类型之间的空格

        SourceLocation spBegLoc = SM.getSpellingLoc(D->getTypeSpecStartLoc());
        const char *begPtr = SM.getCharacterData(spBegLoc);
        
        SourceLocation acEndLoc = Lexer::getLocForEndOfToken(SM.getSpellingLoc(D->getTypeSpecEndLoc()), 0, SM, LangOpts);
        
        const char *endPtr = SM.getCharacterData(acEndLoc);
        
        clang::StringRef fullTypeStr(begPtr, endPtr - begPtr);
        fullTypeStr = fullTypeStr.trim();
        
        if (fullTypeStr.count('*') && !fullTypeStr.count(' ')) {
            DE.Report(acEndLoc.getLocWithOffset(-1), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间至少要有「一个空格」")) << clang::StringRef(begPtr, endPtr - begPtr - 1).trim();
        }
        
        /// 2. 检查名称是否合法
        if (ivarName.size() > 0 && !this->ObjCPropIvarNameCamelCased(ivarName, true)) {
            DE.Report(acEndLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「%0」应以「_ 」开始，并使用「小驼峰命名法」")) << ivarName;
        }
        
        return true;
    }
    
    bool VisitObjCPropertyDecl(const ObjCPropertyDecl *D) {
        if (!this->isUserSourceFile(D->getLocation())) {
            return true;
        }
        
        clang::LangOptions LangOpts = this->Context->getLangOpts();
        clang::TypeLoc loc = D->getTypeSourceInfo()->getTypeLoc();
        clang::StringRef varName = D->getName();
        
        SourceLocation spBegLoc = SM.getSpellingLoc(loc.getBeginLoc());
        const char *begPtr = SM.getCharacterData(spBegLoc);
        
        SourceLocation acEndLoc = Lexer::getLocForEndOfToken(SM.getSpellingLoc(loc.getEndLoc()), 0, SM, LangOpts);
        
        const char *endPtr = SM.getCharacterData(acEndLoc);
        
        clang::StringRef fullTypeStr(begPtr, endPtr - begPtr);
        fullTypeStr = fullTypeStr.trim();
        
        /// 1. 检查修饰符顺序
        {
            SourceLocation lParenLoc = D->getLParenLoc();
            
            const char *lParen = SM.getCharacterData(lParenLoc);
            clang::StringRef propList = clang::StringRef(lParen, begPtr - lParen).trim();
            
            llvm::SmallVector<clang::StringRef> vec;
            propList.substr(1, propList.size() - 2).split(vec, ",", -1, false);
            
            if (!this->ObjCPropertyAttributesOrder(vec)) {
                DE.Report(lParenLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】属性修饰符应遵循以下顺序:\n1.原子修饰符\n2.内存管理修饰符\n3.访问权限修饰符 \n4.可空性修饰符\n5.访问器修饰符\n6.属性类型修饰符"));
            }
        }
        
        /// 2.检查属性命名
        {
            if (!this->ObjCPropIvarNameCamelCased(varName)) {
                DE.Report(D->getBeginLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「%0」应该使用「小驼峰命名法」")) << varName;
            }
        }
        
        clang::QualType type = D->getType();
        clang::ObjCPropertyAttribute::Kind kind = D->getPropertyAttributes();
        
        if (type->isBlockPointerType()) {
            if (!(kind & clang::ObjCPropertyAttribute::kind_copy)) {
                DE.Report(D->getLParenLoc(), DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「block」类型应该使用「copy」修饰，请检查"));
            }
        } else if (type->isPointerType()) {
            /// 1. 检查\*是否与类型中间是否有至少一个空格
            if (!fullTypeStr.count(' ')) {
                DE.Report(acEndLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间至少要有一个空格")) << clang::StringRef(begPtr, endPtr - begPtr - 1).trim();
            }
            
            /// 2. 检查\*与变量名之间是否无空格
            clang::StringRef next = clang::StringRef(endPtr, varName.size());
            if (varName.compare(next)) {
                DE.Report(acEndLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间不应该存在空格")) << varName;
            }
            
            /// 3.检查指针类型
            if (kind & clang::ObjCPropertyAttribute::kind_assign) {
                DE.Report(spBegLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】指针类型应尽量不使用「assign」，请检查"));
            }
        } else if (type->isObjCObjectPointerType()) {
            bool isContainerType = false;
            
            /// 容器类型检查
            {
                if (fullTypeStr.starts_with("NSString") || fullTypeStr.starts_with("NSArray") || fullTypeStr.starts_with("NSDictionary") || fullTypeStr.starts_with("NSSet") || fullTypeStr.starts_with("NSData") || fullTypeStr.starts_with("NSAttributedString") ||
                    fullTypeStr.starts_with("NSIndexSet") ||
                    fullTypeStr.starts_with("NSOrderedSet")) {
                    if (!(kind & clang::ObjCPropertyAttribute::kind_copy)) {
                        DE.Report(spBegLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「不可变容器」类型应尽量使用「copy」修饰符，请检查"));
                    }
                    
                    isContainerType = fullTypeStr.starts_with("NSDictionary") || fullTypeStr.starts_with("NSSet") || fullTypeStr.starts_with("NSArray");
                } else if (fullTypeStr.starts_with("NSMutableString") || fullTypeStr.starts_with("NSMutableArray") || fullTypeStr.starts_with("NSMutableDictionary") || fullTypeStr.starts_with("NSMutableSet") || fullTypeStr.starts_with("NSMutableData") || fullTypeStr.starts_with("NSMutableAttributedString") || fullTypeStr.starts_with("NSMutableIndexSet") || fullTypeStr.starts_with("NSMutableOrderedSet")) {
                    if (kind & clang::ObjCPropertyAttribute::kind_copy) {
                        DE.Report(spBegLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「可变容器」类型不应使用「copy」修饰符，请检查"));
                    }
                    
                    isContainerType = fullTypeStr.starts_with("NSMutableArray") || fullTypeStr.starts_with("NSMutableDictionary") || fullTypeStr.starts_with("NSMutableSet");
                }
            }
            
            /// 泛型容器，必须包含具体类型
            {
                if (isContainerType) {
                    if (fullTypeStr.find("<") == clang::StringRef::npos && fullTypeStr.find(">") == clang::StringRef::npos) {
                        DE.Report(spBegLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【编程规范】「泛型容器」类型应指明具体的类型"));
                    }
                }
            }
            
            {
                /// 检查\*与变量名之间是否无空格
                clang::StringRef next = clang::StringRef(endPtr, varName.size());
                if (varName.compare(next)) {
                    DE.Report(acEndLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「*」与「%0」之间不应该存在空格")) << varName;
                }
                
                /// 3.检查指针类型
                if (kind & clang::ObjCPropertyAttribute::kind_assign) {
                    DE.Report(spBegLoc, DE.getCustomDiagID(this->diagnosticsLevel(), "【代码规范】「指针」类型应尽量不使用「assign」，请检查"));
                }
            }
        }
        
        return true;
    }
    
private:
    clang::CompilerInstance &CI;
    clang::ASTContext *Context;
    SourceManager &SM;
    clang::DiagnosticsEngine &DE;
    std::unique_ptr<clang::ParentMap> PM;
    
    const clang::ObjCMethodDecl *curMethod;
};

class ClangObjCCodeSpecCheckerConsumer : public ASTConsumer {
private:
    CompilerInstance &Instance;
    std::set<std::string> ParsedTemplates;

public:
    explicit ClangObjCCodeSpecCheckerConsumer(CompilerInstance &Instance,
                                              std::set<std::string> ParsedTemplates)
    : Instance(Instance), ParsedTemplates(ParsedTemplates), Visitor(Instance) {}
    
    void HandleTranslationUnit(ASTContext &Ctx) override {
        Visitor.TraverseDecl(Ctx.getTranslationUnitDecl());
    }
    
private:
    ClangOjbCCodeSpecCheckerVisitor Visitor;
};

class ClangObjCCodeSpecCheckerAction : public PluginASTAction {
    std::set<std::string> ParsedTemplates;
protected:
    std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                   llvm::StringRef) override {
        return std::make_unique<ClangObjCCodeSpecCheckerConsumer>(CI, ParsedTemplates);
    }
    
    bool ParseArgs(const CompilerInstance &CI,
                   const std::vector<std::string> &args) override {
        for (auto arg : args) {
            if (!arg.compare(0, 10, "ClsPrefix=")) {
                argsMap["ClsPrefix"] = arg.substr(10, arg.size());
            } else if (!arg.compare(0, 10, "DiagLevel=")) {
                argsMap["DiagLevel"] = arg.substr(10, arg.size());
            } else if (!arg.compare(0, 10, "ArgsCount=")) {
                argsMap["ArgsCount"] = arg.substr(10, arg.size());
            } else if (!arg.compare(0, 13, "MaxLineCount=")) {
                argsMap["MaxLineCount"] = arg.substr(13, arg.size());
            }
        }
        
        return true;
    }
};

}

static FrontendPluginRegistry::Add<ClangObjCCodeSpecCheckerAction>
X("ClangObjCCodeSpecChecker", "clang objective-c code spec checker");
